/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package network;

import network.Layer;

/**
 *
 * @author Raise
 */
import java.util.Random;
import java.util.Arrays;

public class MaxPoolLayer extends Layer {

    private int poolSize;
    private int inputHeight;
    private int inputWidth;
    private int inputChannels;
    private int[][][] maxIndices;

    public MaxPoolLayer(int inputHeight, int inputWidth, int inputChannels, int poolSize) {
        this.inputHeight = inputHeight;
        this.inputWidth = inputWidth;
        this.inputChannels = inputChannels;
        this.poolSize = poolSize;
        this.inputSize = inputHeight * inputWidth * inputChannels;

        int outputHeight = inputHeight / poolSize;
        int outputWidth = inputWidth / poolSize;
        this.outputSize = outputHeight * outputWidth * inputChannels;
    }

    @Override
    public double[] forward(double[] input) {
        double[][][] input3D = reshapeTo3D(input, inputChannels, inputHeight, inputWidth);
        int outputHeight = inputHeight / poolSize;
        int outputWidth = inputWidth / poolSize;

        // 检查整除性
        if (inputHeight % poolSize != 0 || inputWidth % poolSize != 0) {
            throw new IllegalArgumentException("输入尺寸必须能被池化大小整除");
        }

        maxIndices = new int[inputChannels][outputHeight][outputWidth];
        double[] output = new double[outputSize];

        int outputIndex = 0;
        for (int c = 0; c < inputChannels; c++) {
            for (int i = 0; i < outputHeight; i++) {
                for (int j = 0; j < outputWidth; j++) {
                    double maxVal = Double.NEGATIVE_INFINITY;
                    int maxI = -1, maxJ = -1;

                    // 在池化窗口内找最大值
                    for (int pi = 0; pi < poolSize; pi++) {
                        for (int pj = 0; pj < poolSize; pj++) {
                            int inputI = i * poolSize + pi;
                            int inputJ = j * poolSize + pj;
                            double val = input3D[c][inputI][inputJ];
                            if (val > maxVal) {
                                maxVal = val;
                                maxI = inputI;
                                maxJ = inputJ;
                            }
                        }
                    }

                    output[outputIndex++] = maxVal;
                    maxIndices[c][i][j] = maxI * inputWidth + maxJ; // 保存最大值位置
                }
            }
        }

        return output;
    }

    @Override
    public double[] backward(double[] upstreamGradient, double learningRate) {
        // 1. 初始化输入梯度（全0）
        double[] inputGradient = new double[inputSize];
        Arrays.fill(inputGradient, 0.0);

        // 2. 计算输出尺寸
        int outputHeight = inputHeight / poolSize;
        int outputWidth = inputWidth / poolSize;

        // 3. 将上游梯度转为三维
        double[][][] upstream3D = reshapeGradient3D(upstreamGradient, outputHeight, outputWidth, inputChannels);

        // 4. 梯度路由：只传递给最大值位置
        for (int c = 0; c < inputChannels; c++) {
            for (int i = 0; i < outputHeight; i++) {
                for (int j = 0; j < outputWidth; j++) {
                    // 获取前向传播时选中的位置
                    int maxPosition = maxIndices[c][i][j];

                    // 将梯度传递给这个位置
                    inputGradient[maxPosition] += upstream3D[c][i][j];
                }
            }
        }

        return inputGradient;
    }

    private double[][][] reshapeTo3D(double[] input, int channels, int height, int width) {
        double[][][] output = new double[channels][height][width];
        int index = 0;
        for (int c = 0; c < channels; c++) {
            for (int i = 0; i < height; i++) {
                for (int j = 0; j < width; j++) {
                    output[c][i][j] = input[index++];
                }
            }
        }
        return output;
    }

    private double[][][] reshapeGradient3D(double[] flat, int height, int width, int numFilters) {
        double[][][] result = new double[numFilters][height][width];
        int index = 0;
        for (int f = 0; f < numFilters; f++) {
            for (int i = 0; i < height; i++) {
                for (int j = 0; j < width; j++) {
                    result[f][i][j] = flat[index++];
                }
            }
        }
        return result;
    }

   
}
